import os
import json
import re
# import ast
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
from src.llamafactory.chat.chat_model import ChatModel
from tqdm import tqdm
import argparse
import json

model_path = '/root/autodl-tmp/sft_1209_rank64_train0219_5epoch'
parser = argparse.ArgumentParser(description="Your script description")
parser.add_argument("--model_name_or_path", default="/root/autodl-tmp/Qwen2-7B-Instruct", type=str,
                    help="Model name or path")
parser.add_argument("--adapter_name_or_path", default=model_path, type=str, help="Checkpoint directory")
parser.add_argument("--template", default="qwen", type=str, help="Template type")
parser.add_argument("--finetuning_type", default="lora", type=str, help="Fine-tuning type")

args = parser.parse_args()
args_dict = vars(args)
chat_model = ChatModel(args_dict)
history = []

with open('data/test_final.jsonl','r', encoding="utf-8") as f:
    data = f.readlines()

f= open("submission_final.jsonl", "w", encoding="utf-8")

for d in tqdm(data):
    d = json.loads(d)
    apis = d['apis']
    user_message = d['user_messages']

    prompt = f"""你是一名专业的智能体（agent）专家。在智能体执行任务时，需要依据用户多轮对话内容进行总结，明确所需完成的问题，并确定使用哪些工具API以及对应的参数。请按照以下严格要求进行输出：

    1. **API工具识别**：用户的多轮对话可能出现暂时没有工具调用、一个API调用多次、多个API调用或者仅有一次API调用的情况。务必准确给出每个工具的名称，并且确保给出的名称，在提供的API工具中存在。如果用户的一轮对话中调用了多次工具，请分别给出每次调用的工具。
    2. **参数匹配**：严格参照所提供的智能体工具的参数要求，从多轮对话内容中抽取与API所需的必填（required）参数相匹配的信息。不许填写与要求类型不同的参数，所有必填参数都必须依据对话内容给出确切结果。
    3. **必填参数缺失处理**：若必填参数中有任何一项是目前对话无法获取的信息，那么该次对话中不应该给出这个工具的调用。
    4. **非必填参数处理**：对于非必填参数，若对话中未提及或为空，则在参数部分（arguments）中无需列出该参数。
    5. **调用时机判断**：根据当前多轮对话内容，准确判断在哪一轮对话内容能够满足工具调用的参数需求。若在第三轮对话中可明确工具调用的API及其所需参数，则在第三轮给出API调用信息；若在第四轮满足条件，则在第四轮给出。
    6. **多轮对话处理**：请在每一轮对话结束时，整理目前已有的信息给出合理的API调用及参数设置答案。用户有多少轮对话，就返回多少轮。

    请根据以下输入，输出正确结果：
    agents:
    {apis}
    用户多轮对话内容:
    {user_message}
    """
   
    messages = [
        {"role": "user", "content":   prompt}]

    response = chat_model.chat(messages, history=[],  do_sample=False,num_beams=3)
    
    res = response[0].response_text.strip()
    # 把res从字符串转换成list
    try:
        # res = ast.literal_eval(res)
        res = eval(res.replace('true','True').replace('false','False').replace('null','None'))
        assert isinstance(res, list)
    except:
        print(res)
        res = []
    d['targets']=res
    f.write(json.dumps(d, ensure_ascii=False) + "\n")
    f.flush()